import pandas as pd
import numpy as np
from scipy.stats import mannwhitneyu
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
import os
import logging
from datetime import datetime
from typing import Dict, List, Optional, Tuple, Union, Any

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)


class TumorGrowthAnalysis:
    """
    肿瘤生长动力学分析类，使用AUC（曲线下面积）方法评估实验动物的肿瘤生长情况。
    
    该类提供了计算总AUC和分段AUC、处理删失数据、执行统计比较以及生成
    可发表质量图表的工具。
    
    参考文献: 
    Duan, F., et al. (2012). Area under the curve as a tool to
    measure kinetics of tumor growth in experimental animals.
    Journal of Immunological Methods, 382(1-2), 224-228.
    """

    def __init__(self, data: pd.DataFrame, mouse_id_col: str, group_col: str,
                 day_col: str, size_col: str, weight_col: Optional[str] = None):
        """
        使用肿瘤生长数据和可选的体重数据初始化分析对象。

        参数:
            data: 包含肿瘤生长数据的DataFrame。
            mouse_id_col: 包含唯一小鼠标识符的列名。
            group_col: 指定实验组的列名。
            day_col: 包含测量时间点（天）的列名。
            size_col: 包含肿瘤大小测量值的列名。
            weight_col: 包含体重测量值的列名（可选）。
        """
        self._validate_input_data(data, mouse_id_col, group_col, day_col, size_col, weight_col)
        
        # 数据属性初始化
        self.data = data.copy()
        self.mouse_id_col = mouse_id_col
        self.group_col = group_col
        self.day_col = day_col
        self.size_col = size_col
        self.weight_col = weight_col
        
        # 组信息和颜色配置
        self.groups = self.data[self.group_col].unique()
        colors = sns.color_palette("husl", len(self.groups))
        self.color_map = {group: color for group, color in zip(self.groups, colors)}
        self.palette = colors  # 保持向后兼容
        
        # 结果存储属性
        self.auc_results = None  # 向后兼容
        self.segmented_auc_results = None  # 向后兼容
        self.tumor_auc_results = None
        self.tumor_segmented_auc_results = None
        self.weight_auc_results = None
        self.weight_segmented_auc_results = None
        self.weight_change_results = None
        self.weight_rmsd_results = None
    
    def _validate_input_data(self, data: pd.DataFrame, mouse_id_col: str,
                            group_col: str, day_col: str, size_col: str,
                            weight_col: Optional[str]) -> None:
        """
        验证输入数据的有效性。
        
        参数:
            data: 输入的DataFrame。
            mouse_id_col: 小鼠ID列名。
            group_col: 组列名。
            day_col: 天数列名。
            size_col: 肿瘤大小列名。
            weight_col: 体重列名（可选）。
        
        异常:
            ValueError: 当必要的列不存在时抛出。
        """
        required_cols = [mouse_id_col, group_col, day_col, size_col]
        if not all(col in data.columns for col in required_cols):
            missing_cols = [col for col in required_cols if col not in data.columns]
            raise ValueError(f"数据缺少必要的列: {missing_cols}")

        if weight_col and weight_col not in data.columns:
            raise ValueError(f"体重列 '{weight_col}' 不存在于数据中")
    
    def _calculate_single_auc(self, mouse_df: pd.DataFrame, total_duration: Union[int, float], 
                             metric_col: str) -> float:
        """
        计算单个小鼠的AUC，处理删失数据。
        
        参数:
            mouse_df: 单个小鼠的数据。
            total_duration: 实验的总持续时间。
            metric_col: 用于计算AUC的指标列名。
            
        返回:
            计算得到的AUC值。
        """
        mouse_df = mouse_df.sort_values(by=self.day_col).reset_index(drop=True)
        
        days = mouse_df[self.day_col].values
        values = mouse_df[metric_col].values
        
        last_day = days[-1]
        
        # 通过水平延伸处理删失数据
        if last_day < total_duration:
            days = np.append(days, total_duration)
            values = np.append(values, values[-1])  # 用最后已知值延伸
            
        # 使用np.trapezoid计算曲线下面积
        return np.trapezoid(values, days)

    def calculate_auc(self, total_duration: Union[int, float], metric: str = 'tumor') -> pd.DataFrame:
        """
        计算指定持续时间内每只小鼠的总AUC。

        参数:
            total_duration: 实验的总持续时间，用于纠正删失数据。
            metric: 要分析的指标，'tumor'表示肿瘤体积，'weight'表示体重。

        返回:
            包含mouse_id、group和计算出的total_auc的DataFrame。
        """
        metric_col = self._get_metric_column(metric)
        
        auc_data = []
        for mouse_id, group_df in self.data.groupby(self.mouse_id_col):
            try:
                auc = self._calculate_single_auc(group_df, total_duration, metric_col)
                group = group_df[self.group_col].iloc[0]
                auc_data.append({
                    self.mouse_id_col: mouse_id,
                    self.group_col: group,
                    'total_auc': auc
                })
            except Exception as e:
                logger.warning(f"计算小鼠 {mouse_id} 的AUC时出错: {e}")
                continue
        
        if not auc_data:
            raise ValueError(f"无法计算任何小鼠的{metric} AUC")
        
        result_df = pd.DataFrame(auc_data)
        
        # 根据指标类型存储结果
        if metric == 'tumor':
            self.tumor_auc_results = result_df
        else:
            self.weight_auc_results = result_df
        self.auc_results = result_df  # 保持向后兼容
        
        return result_df
    
    def _get_metric_column(self, metric: str) -> str:
        """
        根据指标类型获取对应的列名。
        
        参数:
            metric: 指标类型，'tumor'或'weight'。
            
        返回:
            对应的列名。
        
        异常:
            ValueError: 当指标无效或体重列未提供时抛出。
        """
        if metric == 'tumor':
            return self.size_col
        elif metric == 'weight':
            if not self.weight_col:
                raise ValueError("未提供体重列。请在初始化时指定weight_col参数。")
            return self.weight_col
        else:
            raise ValueError("无效的指标。请选择 'tumor' 或 'weight'。")
    
    def _calculate_single_segmented_auc(self, mouse_df: pd.DataFrame, segments: List[Union[int, float]],
                                       metric_col: str) -> Dict[str, float]:
        """
        计算单个小鼠的分段AUC。
        
        参数:
            mouse_df: 单个小鼠的数据。
            segments: 时间分段列表。
            metric_col: 用于计算AUC的指标列名。
            
        返回:
            包含各段AUC值的字典。
        """
        mouse_df = mouse_df.sort_values(by=self.day_col).reset_index(drop=True)
        days = mouse_df[self.day_col].values
        values = mouse_df[metric_col].values
        
        aucs = {}
        for i in range(len(segments) - 1):
            start_day, end_day = segments[i], segments[i+1]
            segment_name = f'auc_{start_day}-{end_day}'
            
            # 找到段内的点并插值
            segment_mask = (days >= start_day) & (days <= end_day)
            interp_days = np.unique(np.clip([start_day, end_day] + list(days[segment_mask]), start_day, end_day))
            interp_values = np.interp(interp_days, days, values)
            
            aucs[segment_name] = np.trapezoid(interp_values, interp_days)
            
        return aucs

    def calculate_segmented_auc(self, segments: List[Union[int, float]], metric: str = 'tumor') -> pd.DataFrame:
        """
        计算预定义时间段内每只小鼠的AUC。

        参数:
            segments: 定义段的时间点列表（例如，[0, 10, 20, 30]）。
            metric: 要分析的指标，'tumor'表示肿瘤体积，'weight'表示体重。

        返回:
            包含mouse_id、group和每个段AUC的DataFrame。
        """
        if len(segments) < 2:
            raise ValueError("段必须包含至少两个时间点。")
            
        metric_col = self._get_metric_column(metric)
            
        segmented_data = []
        for mouse_id, group_df in self.data.groupby(self.mouse_id_col):
            try:
                aucs = self._calculate_single_segmented_auc(group_df, segments, metric_col)
                group = group_df[self.group_col].iloc[0]
                record = {self.mouse_id_col: mouse_id, self.group_col: group, **aucs}
                segmented_data.append(record)
            except Exception as e:
                logger.warning(f"计算小鼠 {mouse_id} 的分段AUC时出错: {e}")
                continue
                
        if not segmented_data:
            raise ValueError(f"无法计算任何小鼠的{metric}分段AUC")
            
        result_df = pd.DataFrame(segmented_data)
        
        # 根据指标类型存储结果
        if metric == 'tumor':
            self.tumor_segmented_auc_results = result_df
        else:
            self.weight_segmented_auc_results = result_df
        self.segmented_auc_results = result_df  # 保持向后兼容
        
        return result_df

    def compare_groups(self, metric: str, group1: str, group2: str, test: str = 'mann-whitneyu') -> Dict[str, Any]:
        """
        执行统计测试以比较两组之间的指标。

        参数:
            metric: 要比较的指标名称（例如，'total_auc'或'auc_0-10'）。
            group1: 第一组的名称。
            group2: 第二组的名称。
            test: 要使用的统计测试。目前仅支持'mann-whitneyu'。

        返回:
            包含测试结果（p值）的字典。
        """
        data_source = self._get_metric_data_source(metric)

        vals1 = data_source[data_source[self.group_col] == group1][metric].dropna()
        vals2 = data_source[data_source[self.group_col] == group2][metric].dropna()

        if len(vals1) < 1 or len(vals2) < 1:
            logger.warning(f"比较 {group1} 和 {group2} 时数据不足")
            return {'p_value': np.nan, 'statistic': np.nan, 'note': '比较数据不足'}

        if test == 'mann-whitneyu':
            stat, p_val = mannwhitneyu(vals1, vals2, alternative='two-sided')
            return {'p_value': p_val, 'statistic': stat}
        else:
            raise NotImplementedError(f"测试 '{test}' 尚未实现。")
    
    def _get_metric_data_source(self, metric: str) -> pd.DataFrame:
        """
        根据指标名称获取相应的数据来源。
        
        参数:
            metric: 指标名称。
            
        返回:
            包含该指标数据的DataFrame。
            
        异常:
            ValueError: 当未找到相应的数据时抛出。
        """
        if metric == 'total_auc':
            if self.auc_results is not None:
                return self.auc_results
            else:
                raise ValueError("未计算总AUC。请先运行calculate_auc。")
        elif metric.startswith('auc_'):
            if self.segmented_auc_results is not None:
                return self.segmented_auc_results
            else:
                raise ValueError("未计算分段AUC。请先运行calculate_segmented_auc。")
        else:
            raise ValueError(f"未找到指标 '{metric}'。请先运行calculate_auc或calculate_segmented_auc。")

    def plot_growth_curves(self, metric: str = 'tumor', ax: Optional[plt.Axes] = None,
                          save_path: Optional[str] = None, **kwargs) -> plt.Axes:
        """
        绘制每组的平均肿瘤生长曲线或体重变化曲线。

        参数:
            metric: 要绘制的指标，'tumor'表示肿瘤体积，'weight'表示体重。
            ax: 用于绘图的现有轴。
            save_path: 保存图形的路径。
            **kwargs: 传递给seaborn.lineplot的其他关键字参数。
            
        返回:
            matplotlib轴对象。
        """
        y_col, title, ylabel = self._get_plot_parameters(metric)
        
        if ax is None:
            fig, ax = plt.subplots(figsize=(8, 6))
        
        sns.lineplot(
            data=self.data,
            x=self.day_col,
            y=y_col,
            hue=self.group_col,
            errorbar='se',  # 标准误差
            ax=ax,
            palette=self.color_map,
            **kwargs
        )
        
        ax.set_title(title)
        ax.set_xlabel('Days Post Challenge')
        ax.set_ylabel(ylabel)
        ax.legend(title='Group')
        
        if save_path:
            self._save_figure(save_path)
            
        return ax
    
    def _get_plot_parameters(self, metric: str) -> Tuple[str, str, str]:
        """
        根据指标类型获取绘图参数。
        
        参数:
            metric: 指标类型，'tumor'或'weight'。
            
        返回:
            包含y列名、图表标题和y轴标签的元组。
        """
        if metric == 'tumor':
            return self.size_col, 'Tumor Growth Curves', f'Tumor Size ({self.size_col})'
        elif metric == 'weight':
            if not self.weight_col:
                raise ValueError("未提供体重列。请在初始化时指定weight_col参数。")
            return self.weight_col, 'Body Weight Change Curves', f'Body Weight ({self.weight_col})'
        else:
            raise ValueError("无效的指标。请选择 'tumor' 或 'weight'。")
    
    def _save_figure(self, save_path: str) -> None:
        """
        保存图形到指定路径。
        
        参数:
            save_path: 保存路径。
        """
        try:
            plt.savefig(save_path, bbox_inches='tight', dpi=300)
            logger.info(f"图形已保存到: {save_path}")
        except Exception as e:
            logger.error(f"保存图形失败: {e}")
    
    def _create_plot(self, data: pd.DataFrame, x_col: str, y_col: str, plot_type: str,
                    ax: plt.Axes, **kwargs) -> plt.Axes:
        """
        创建指定类型的图表。
        
        参数:
            data: 绘图数据。
            x_col: x轴列名。
            y_col: y轴列名。
            plot_type: 图表类型。
            ax: matplotlib轴对象。
            **kwargs: 传递给绘图函数的其他参数。
            
        返回:
            matplotlib轴对象。
        """
        if plot_type == 'box':
            sns.boxplot(
                data=data, 
                x=x_col, 
                y=y_col, 
                hue=x_col, 
                dodge=False, 
                ax=ax, 
                palette=self.color_map, 
                legend=False,
                **kwargs
            )
            sns.stripplot(data=data, x=x_col, y=y_col, ax=ax, color='black', alpha=0.5)
        elif plot_type == 'bar':
            sns.barplot(
                data=data, 
                x=x_col, 
                y=y_col, 
                hue=x_col,
                dodge=False,
                errorbar='sd', 
                ax=ax, 
                palette=self.color_map,
                legend=False,
                **kwargs
            )
            sns.stripplot(data=data, x=x_col, y=y_col, ax=ax, color='black', alpha=0.5)
        elif plot_type == 'strip':
            sns.stripplot(
                data=data, 
                x=x_col, 
                y=y_col, 
                hue=x_col,
                ax=ax, 
                palette=self.color_map,
                legend=False,
                **kwargs
            )
        elif plot_type == 'violin':
            sns.violinplot(
                data=data, 
                x=x_col, 
                y=y_col, 
                hue=x_col,
                dodge=False,
                ax=ax, 
                palette=self.color_map,
                legend=False,
                **kwargs
            )
            sns.stripplot(data=data, x=x_col, y=y_col, ax=ax, color='black', alpha=0.5)
        elif plot_type == 'swarm':
            sns.swarmplot(
                data=data, 
                x=x_col, 
                y=y_col, 
                hue=x_col,
                ax=ax, 
                palette=self.color_map,
                legend=False,
                **kwargs
            )
        else:
            raise ValueError(f"不支持的图表类型: {plot_type}。支持的类型包括: 'box', 'bar', 'strip', 'violin', 'swarm'")
        
        return ax

    def plot_auc_summary(self, metric: str, control_group: Optional[str] = None,
                        plot_type: str = 'box', ax: Optional[plt.Axes] = None,
                        save_path: Optional[str] = None, **kwargs) -> plt.Axes:
        """
        为所有组生成AUC指标的汇总图。

        参数:
            metric: 要绘制的AUC指标。
            control_group: 要比较的对照组。如果提供，将添加统计注释。
            plot_type: 图类型（'box', 'bar', 'strip'等）。
            ax: 用于绘图的现有轴。
            save_path: 保存图形的路径。
            **kwargs: 传递给seaborn绘图函数的其他关键字参数。
            
        返回:
            matplotlib轴对象。
        """
        data_source = self._get_metric_data_source(metric)

        if ax is None:
            fig, ax = plt.subplots(figsize=(7, 6))
        
        # 创建基础图表
        ax = self._create_plot(data_source, self.group_col, metric, plot_type, ax, **kwargs)

        # 设置标题和标签
        title, ylabel = self._get_auc_plot_title(metric)
        ax.set_title(title)
        ax.set_xlabel('Group')
        ax.set_ylabel(ylabel)
        
        # 添加统计注释
        if control_group:
            self._add_statistical_annotations(ax, data_source, metric, control_group)

        if save_path:
            self._save_figure(save_path)
            
        return ax
    
    def _get_auc_plot_title(self, metric: str) -> Tuple[str, str]:
        """
        获取AUC图表的标题和y轴标签。
        
        参数:
            metric: AUC指标名称。
            
        返回:
            包含标题和y轴标签的元组。
        """
        if metric == 'total_auc':
            if self.weight_auc_results is not None:
                # 判断当前使用的是肿瘤还是体重数据
                if 'weight_g' in self.data.columns:
                    return 'Comparison of Total Weight AUC', 'Total Weight AUC'
                else:
                    return 'Comparison of Total Tumor AUC', 'Total Tumor AUC'
            else:
                return 'Comparison of Total Tumor AUC', 'Total Tumor AUC'
        else:
            # 分段AUC标题
            if self.weight_segmented_auc_results is not None:
                if 'weight_g' in self.data.columns:
                    return f'Comparison of Weight {metric}', f'Weight {metric}'
                else:
                    return f'Comparison of Tumor {metric}', f'Tumor {metric}'
            else:
                return f'Comparison of Tumor {metric}', f'Tumor {metric}'
    
    def _add_statistical_annotations(self, ax: plt.Axes, data: pd.DataFrame, 
                                    metric: str, control_group: str) -> None:
        """
        添加统计显著性注释。
        
        参数:
            ax: matplotlib轴对象。
            data: 包含数据的DataFrame。
            metric: 要比较的指标。
            control_group: 对照组名称。
        """
        # 获取所有唯一的组名
        all_groups = data[self.group_col].unique()
        logger.info(f"所有组名: {all_groups}")
        logger.info(f"指定的对照组: {control_group}")
        
        # 确保对照组存在
        if control_group not in all_groups:
            logger.warning(f"对照组 '{control_group}' 不存在于数据中")
            return
        
        # 获取所有非对照组
        experimental_groups = [g for g in all_groups if g != control_group]
        logger.info(f"实验组列表: {experimental_groups}")
        
        # 设置显著性线的位置参数
        y_max = data[metric].max()
        y_step = y_max * 0.1
        y_pos = y_max + y_step

        # 为每个实验组与对照组进行比较
        for i, exp_group in enumerate(experimental_groups):
            try:
                logger.info(f"执行统计检验: {control_group} vs {exp_group}, 指标: {metric}")
                # 执行对照组与实验组的比较
                res = self.compare_groups(metric, control_group, exp_group)
                p_val = res['p_value']
                logger.info(f"{control_group} vs {exp_group} 的p值: {p_val}")
                
                # 确定显著性标记
                if p_val < 0.001:
                    p_text = '***'
                elif p_val < 0.01:
                    p_text = '**'
                elif p_val < 0.05:
                    p_text = '*'
                else:
                    p_text = 'ns'  # 无显著性差异
                
                # 直接从ax获取组顺序，避免重复创建boxplot
                group_order = [tick.get_text() for tick in ax.get_xticklabels()]
                logger.info(f"图表中的组顺序: {group_order}")
                
                # 确保组名在顺序中
                if control_group in group_order and exp_group in group_order:
                    x1 = group_order.index(control_group)
                    x2 = group_order.index(exp_group)
                    logger.info(f"{control_group}的x坐标: {x1}, {exp_group}的x坐标: {x2}")
                    
                    # 绘制显著性线和文本
                    ax.plot([x1, x1, x2, x2], [y_pos, y_pos + y_step*0.2, y_pos + y_step*0.2, y_pos], 
                           lw=1.5, c='black')
                    ax.text((x1 + x2) * 0.5, y_pos + y_step*0.2, p_text, ha='center', va='bottom', color='black')
                    logger.info(f"添加统计注释: {control_group} vs {exp_group}, p={p_val:.6f}, 显著性={p_text}")
                    
                    # 更新y_pos为下一个比较
                    y_pos += y_step
                else:
                    logger.warning(f"组 {control_group} 或 {exp_group} 不在图表中")
            except Exception as e:
                logger.error(f"计算 {control_group} 与 {exp_group} 的统计显著性时出错: {str(e)}")
                logger.warning(f"添加 {control_group} 和 {exp_group} 之间的统计注释时出错: {e}")
                continue

    def calculate_weight_change(self) -> pd.DataFrame:
        """
        计算每只小鼠的体重变化百分比和最终体重。
        
        返回:
            包含mouse_id、group、初始体重、最终体重和体重变化百分比的DataFrame。
        """
        if not self.weight_col:
            raise ValueError("未指定体重列。请在初始化时提供weight_col参数。")
        
        weight_data = []
        
        # 过滤掉无效数据
        valid_data = self._filter_valid_data(self.weight_col)
        
        for mouse_id, mouse_df in valid_data.groupby(self.mouse_id_col):
            try:
                mouse_df = mouse_df.sort_values(by=self.day_col)
                
                # 获取初始体重和最终体重
                initial_weight = mouse_df.iloc[0][self.weight_col]  # 第一天的数据
                final_weight = mouse_df.iloc[-1][self.weight_col]  # 最后一天的数据
                
                # 计算体重变化百分比
                weight_change_pct = ((final_weight - initial_weight) / initial_weight) * 100
                
                group = mouse_df.iloc[0][self.group_col]
                
                weight_data.append({
                    self.mouse_id_col: mouse_id,
                    self.group_col: group,
                    'initial_weight': initial_weight,
                    'final_weight': final_weight,
                    'weight_change_pct': weight_change_pct
                })
            except Exception as e:
                logger.warning(f"处理小鼠 {mouse_id} 时出错: {e}")
                continue
        
        if not weight_data:
            raise ValueError("没有找到有效的体重数据进行分析")
            
        self.weight_change_results = pd.DataFrame(weight_data)
        return self.weight_change_results
    
    def _filter_valid_data(self, value_col: str) -> pd.DataFrame:
        """
        过滤有效的数据行。
        
        参数:
            value_col: 要包含的数值列。
            
        返回:
            过滤后的DataFrame。
        """
        valid_data = self.data.dropna(subset=[self.mouse_id_col, self.group_col, self.day_col, value_col])
        # 过滤掉空字符串组名
        valid_data = valid_data[valid_data[self.group_col].str.strip().astype(bool)]
        return valid_data
    
    def plot_weight_curves(self, ax: Optional[plt.Axes] = None, 
                          save_path: Optional[str] = None, **kwargs) -> plt.Axes:
        """
        绘制每组的平均体重变化曲线。
        
        参数:
            ax: 用于绘图的现有轴。
            save_path: 保存图形的路径。
            **kwargs: 传递给seaborn.lineplot的其他关键字参数。
            
        返回:
            matplotlib轴对象。
        """
        if not self.weight_col:
            raise ValueError("未指定体重列。请在初始化时提供weight_col参数。")
        
        # 复用plot_growth_curves方法
        return self.plot_growth_curves(metric='weight', ax=ax, save_path=save_path, **kwargs)
    
    def calculate_weight_rmsd(self) -> Dict[str, float]:
        """
        计算每只小鼠的体重RMSD（均方根偏差）和每个组的平均体重RMSD。
        
        返回:
            包含每个组平均体重RMSD的字典。
        """
        if not self.weight_col:
            raise ValueError("未指定体重列。请在初始化时提供weight_col参数。")
        
        valid_data = self._filter_valid_data(self.weight_col)
        
        weight_rmsd_dict = {}
        
        for group in valid_data[self.group_col].unique():
            try:
                group_data = valid_data[valid_data[self.group_col] == group]
                mouse_ids = group_data[self.mouse_id_col].unique()
                mouse_rmsds = []
                
                for mouse_id in mouse_ids:
                    mouse_data = group_data[group_data[self.mouse_id_col] == mouse_id].sort_values(self.day_col)
                    weights = mouse_data[self.weight_col].values
                    avg_weight = np.mean(weights)
                    rmsd = np.sqrt(np.mean((weights - avg_weight) ** 2))
                    mouse_rmsds.append(rmsd)
                
                # 计算该组的平均RMSD
                group_avg_rmsd = np.mean(mouse_rmsds)
                weight_rmsd_dict[group] = group_avg_rmsd
            except Exception as e:
                logger.warning(f"计算组 {group} 的体重RMSD时出错: {e}")
                continue
        
        if not weight_rmsd_dict:
            raise ValueError("无法计算任何组的体重RMSD")
            
        self.weight_rmsd_results = weight_rmsd_dict
        return weight_rmsd_dict
    
    def plot_weight_change_summary(self, plot_type: str = 'box', 
                                 ax: Optional[plt.Axes] = None,
                                 save_path: Optional[str] = None, **kwargs) -> plt.Axes:
        """
        为所有组生成体重变化百分比的汇总图。
        
        参数:
            plot_type: 图类型（'box', 'bar', 'strip'等）。
            ax: 用于绘图的现有轴。
            save_path: 保存图形的路径。
            **kwargs: 传递给seaborn绘图函数的其他关键字参数。
            
        返回:
            matplotlib轴对象。
        """
        if self.weight_change_results is None or self.weight_change_results.empty:
            self.calculate_weight_change()
        
        if ax is None:
            fig, ax = plt.subplots(figsize=(7, 6))
        
        # 创建基础图表
        ax = self._create_plot(self.weight_change_results, self.group_col, 
                              'weight_change_pct', plot_type, ax, **kwargs)

        # 设置标题和标签
        ax.set_title('Comparison of Weight Change Percentage')
        ax.set_xlabel('Group')
        ax.set_ylabel('Weight Change (%)')
        
        # 添加零线以更容易观察体重增加和减少
        ax.axhline(y=0, color='r', linestyle='--', alpha=0.3)
        
        if save_path:
            self._save_figure(save_path)
            
        return ax


def set_palette(analysis: TumorGrowthAnalysis, palette_name: str, group_count: int) -> None:
    """
    设置分析对象的颜色映射。
    
    参数:
        analysis: TumorGrowthAnalysis对象。
        palette_name: 调色板名称或'custom'。
        group_count: 组的数量。
    """
    if palette_name == 'custom':
        # 自定义颜色列表
        custom_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', 
                        '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
        colors = custom_colors[:group_count]
    else:
        # 使用seaborn内置调色板
        colors = sns.color_palette(palette_name, group_count)
    
    # 创建组名到颜色的映射
    analysis.color_map = {group: color for group, color in zip(analysis.groups, colors)}
    # 更新palette属性以保持向后兼容
    analysis.palette = colors


class TumorAUCCommand:
    """
    肿瘤AUC分析命令行工具类。
    """
    
    def __init__(self):
        """
        初始化命令行参数。
        """
        self.parser = argparse.ArgumentParser(description='使用AUC方法进行肿瘤生长和体重变化分析')
        self._setup_arguments()
    
    def _setup_arguments(self) -> None:
        """
        设置命令行参数。
        """
        self.parser.add_argument('--input', type=str, default='sample_tumor_data.csv', 
                               help='包含肿瘤生长数据的输入CSV文件路径')
        self.parser.add_argument('--seg', type=int, default=3, 
                               help='分段数目，用于将实验期平均分成若干段')
        self.parser.add_argument('--fig', type=str, default='box', 
                               help='绘图类型，可选值: box(箱型图), bar(柱状图), strip(散点图), violin(小提琴图), swarm(蜂群图)')
        self.parser.add_argument('--palette', type=str, default='husl', 
                               help='颜色方案，可选值: husl, tab10, Set2, viridis, plasma等')
        self.parser.add_argument('--metric', type=str, default='tumor', 
                               help='分析指标，可选值: tumor(肿瘤体积), weight(体重)')
    
    def run(self) -> None:
        """
        运行肿瘤AUC分析工作流。
        """
        args = self.parser.parse_args()
        
        # 创建结果存储目录（使用时间戳）
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        result_dir = os.path.join('result', timestamp)
        os.makedirs(result_dir, exist_ok=True)
        logger.info(f"结果将保存到: {result_dir}")

        # 1. 从CSV文件加载数据
        try:
            df = pd.read_csv(args.input)
        except FileNotFoundError:
            logger.error(f"错误: '{args.input}' 文件未找到。请确保该文件存在。")
            return
        except Exception as e:
            logger.error(f"读取CSV文件时出错: {e}")
            return

        # 2. 初始化分析对象
        try:
            analysis = TumorGrowthAnalysis(
                data=df,
                mouse_id_col='mouse_id',
                group_col='group',
                day_col='day',
                size_col='tumor_volume_mm3',
                weight_col='weight_g'  # 添加体重列支持
            )
        except Exception as e:
            logger.error(f"初始化分析对象时出错: {e}")
            return
        
        # 3. 设置颜色方案
        group_count = len(analysis.groups)
        set_palette(analysis, args.palette, group_count)

        # 4. 绘制原始生长曲线
        try:
            self._generate_growth_curves(analysis, args.metric, result_dir)
        except Exception as e:
            logger.error(f"生成生长曲线时出错: {e}")

        # 5. 计算总AUC
        try:
            total_auc_results = self._calculate_total_auc(analysis, args.metric)
        except Exception as e:
            logger.error(f"计算总AUC时出错: {e}")
            total_auc_results = None

        # 6. 绘制总AUC汇总图
        if total_auc_results is not None:
            try:
                self._plot_total_auc_summary(analysis, args.metric, args.fig, result_dir)
            except Exception as e:
                logger.error(f"绘制总AUC汇总图时出错: {e}")

        # 7. 计算分段AUC
        try:
            segmented_auc_results = self._calculate_segmented_auc(analysis, args.metric, args.seg)
        except Exception as e:
            logger.error(f"计算分段AUC时出错: {e}")
            segmented_auc_results = None

        # 8. 为每个分段生成AUC比较图
        if segmented_auc_results is not None and total_auc_results is not None:
            try:
                self._generate_segmented_auc_plots(analysis, args.metric, args.fig, args.seg, result_dir)
                
                # 9. 生成并保存小鼠个体AUC数据CSV文件
                self._save_individual_auc_data(total_auc_results, segmented_auc_results, 
                                              analysis, result_dir)
                
                # 10. 计算并保存小组平均AUC数据CSV文件
                self._save_group_avg_auc_data(total_auc_results, segmented_auc_results, 
                                            analysis, result_dir)
            except Exception as e:
                logger.error(f"处理分段AUC数据时出错: {e}")

        # 11-15. 分析体重变化数据
        try:
            self._analyze_weight_data(analysis, args.fig, result_dir)
        except Exception as e:
            logger.error(f"分析体重数据时出错: {e}")

        # 16. 执行特定组的比较
        try:
            self._perform_specific_comparison(analysis, args.metric)
        except Exception as e:
            logger.error(f"执行特定比较时出错: {e}")
    
    def _generate_growth_curves(self, analysis: TumorGrowthAnalysis, metric: str, 
                              result_dir: str) -> None:
        """
        生成生长曲线。
        """
        if metric == 'tumor':
            logger.info("Generating tumor growth curves plot...")
            growth_curves_path = os.path.join(result_dir, 'tumor_growth_curves.png')
        else:
            logger.info("Generating body weight change curves plot...")
            growth_curves_path = os.path.join(result_dir, 'weight_change_curves.png')
        
        analysis.plot_growth_curves(metric=metric, save_path=growth_curves_path)
        plt.close()  # 关闭图形以释放内存
    
    def _calculate_total_auc(self, analysis: TumorGrowthAnalysis, metric: str) -> pd.DataFrame:
        """
        计算总AUC。
        """
        logger.info(f"\nCalculating total {metric} AUC for each mouse (total duration = 28 days)...")
        total_auc_results = analysis.calculate_auc(total_duration=28, metric=metric)
        logger.info(f"Total AUC results:\n{total_auc_results}")
        return total_auc_results
    
    def _plot_total_auc_summary(self, analysis: TumorGrowthAnalysis, metric: str, 
                              plot_type: str, result_dir: str) -> None:
        """
        绘制总AUC汇总图。
        """
        logger.info(f"\nGenerating total {metric} AUC summary plot with statistics...")
        if metric == 'tumor':
            total_auc_plot_path = os.path.join(result_dir, f'total_tumor_auc_comparison_{plot_type}.png')
        else:
            total_auc_plot_path = os.path.join(result_dir, f'total_weight_auc_comparison_{plot_type}.png')
        
        analysis.plot_auc_summary(
            metric='total_auc', 
            control_group='Control',
            plot_type=plot_type,
            save_path=total_auc_plot_path
        )
        plt.close()
    
    def _calculate_segmented_auc(self, analysis: TumorGrowthAnalysis, metric: str, 
                               seg: int) -> pd.DataFrame:
        """
        计算分段AUC。
        """
        total_duration = 28
        segments = list(np.linspace(0, total_duration, seg + 1, dtype=int))
        logger.info(f"\nCalculating segmented {metric} AUC for periods: {segments}...")
        
        segmented_auc_results = analysis.calculate_segmented_auc(segments=segments, metric=metric)
        logger.info(f"Segmented AUC results:\n{segmented_auc_results}")
        return segmented_auc_results
    
    def _generate_segmented_auc_plots(self, analysis: TumorGrowthAnalysis, metric: str, 
                                    plot_type: str, seg: int, result_dir: str) -> None:
        """
        生成分段AUC比较图。
        """
        logger.info(f"\nGenerating {metric} AUC comparison plots for all segments...")
        total_duration = 28
        segments = list(np.linspace(0, total_duration, seg + 1, dtype=int))
        
        for i in range(len(segments) - 1):
            start_day, end_day = segments[i], segments[i+1]
            segment_metric = f'auc_{start_day}-{end_day}'
            
            if metric == 'tumor':
                save_path = os.path.join(result_dir, f'segmented_tumor_auc_{start_day}-{end_day}_comparison_{plot_type}.png')
            else:
                save_path = os.path.join(result_dir, f'segmented_weight_auc_{start_day}-{end_day}_comparison_{plot_type}.png')

            logger.info(f"  Generating plot for {segment_metric}...")
            analysis.plot_auc_summary(
                metric=segment_metric,
                control_group='Control',
                plot_type=plot_type,
                save_path=save_path
            )
            plt.close()
    
    def _save_individual_auc_data(self, total_auc_results: pd.DataFrame, 
                                segmented_auc_results: pd.DataFrame, 
                                analysis: TumorGrowthAnalysis, result_dir: str) -> None:
        """
        保存个体小鼠AUC数据。
        """
        logger.info("\nGenerating individual mouse AUC data CSV...")
        # 合并total_auc和segmented_auc数据
        merged_auc_data = pd.merge(
            total_auc_results,
            segmented_auc_results,
            on=[analysis.mouse_id_col, analysis.group_col],
            how='outer'
        )
        # 保存到CSV
        individual_auc_csv_path = os.path.join(result_dir, 'individual_mouse_auc_data.csv')
        merged_auc_data.to_csv(individual_auc_csv_path, index=False)
        logger.info(f"个体小鼠AUC数据已保存到: {individual_auc_csv_path}")
    
    def _save_group_avg_auc_data(self, total_auc_results: pd.DataFrame, 
                                segmented_auc_results: pd.DataFrame, 
                                analysis: TumorGrowthAnalysis, result_dir: str) -> None:
        """
        保存小组平均AUC数据。
        """
        logger.info("\nGenerating group average AUC data CSV...")
        # 合并数据以获取所有列
        merged_auc_data = pd.merge(
            total_auc_results,
            segmented_auc_results,
            on=[analysis.mouse_id_col, analysis.group_col],
            how='outer'
        )
        
        # 获取所有AUC列
        auc_columns = [col for col in merged_auc_data.columns if col.endswith('_auc') or col.startswith('auc_')]
        # 按组计算平均值
        group_avg_auc = merged_auc_data.groupby(analysis.group_col)[auc_columns].mean().reset_index()
        # 重命名列以指示平均值
        group_avg_auc.columns = [f"{col}_mean" if col in auc_columns else col for col in group_avg_auc.columns]
        # 保存到CSV
        group_avg_auc_csv_path = os.path.join(result_dir, 'group_average_auc_data.csv')
        group_avg_auc.to_csv(group_avg_auc_csv_path, index=False)
        logger.info(f"小组平均AUC数据已保存到: {group_avg_auc_csv_path}")
    
    def _analyze_weight_data(self, analysis: TumorGrowthAnalysis, plot_type: str, 
                           result_dir: str) -> None:
        """
        分析体重变化数据。
        """
        logger.info("\nAnalyzing body weight changes...")
        weight_change_results = analysis.calculate_weight_change()
        logger.info(f"Weight change results:\n{weight_change_results}")
        
        # 绘制体重变化曲线
        logger.info("\nGenerating body weight curves plot...")
        weight_curves_path = os.path.join(result_dir, 'body_weight_curves.png')
        analysis.plot_weight_curves(save_path=weight_curves_path)
        plt.close()
        
        # 绘制体重变化百分比汇总图
        logger.info("\nGenerating weight change percentage summary plot...")
        weight_change_plot_path = os.path.join(result_dir, f'weight_change_comparison_{plot_type}.png')
        analysis.plot_weight_change_summary(
            plot_type=plot_type,
            save_path=weight_change_plot_path
        )
        plt.close()
        
        # 保存体重变化数据到CSV
        logger.info("\nGenerating individual mouse weight change data CSV...")
        weight_change_csv_path = os.path.join(result_dir, 'individual_mouse_weight_change.csv')
        weight_change_results.to_csv(weight_change_csv_path, index=False)
        logger.info(f"个体小鼠体重变化数据已保存到: {weight_change_csv_path}")
        
        # 计算并保存小组平均体重变化数据CSV
        logger.info("\nGenerating group average weight change data CSV...")
        weight_columns = ['initial_weight', 'final_weight', 'weight_change_pct']
        group_avg_weight = weight_change_results.groupby(analysis.group_col)[weight_columns].mean().reset_index()
        group_avg_weight.columns = [f"{col}_mean" if col in weight_columns else col for col in group_avg_weight.columns]
        
        # 添加体重RMSD计算
        logger.info("\nCalculating weight RMSD for each group...")
        weight_rmsd_results = analysis.calculate_weight_rmsd()
        # 将RMSD结果添加到group_avg_weight DataFrame
        group_avg_weight['weight_rmsd'] = group_avg_weight[analysis.group_col].map(weight_rmsd_results)
        logger.info("Weight RMSD results:")
        for group, rmsd in weight_rmsd_results.items():
            logger.info(f"  {group}: {rmsd:.4f}")
        
        group_avg_weight_csv_path = os.path.join(result_dir, 'group_average_weight_change.csv')
        group_avg_weight.to_csv(group_avg_weight_csv_path, index=False)
        logger.info(f"小组平均体重变化数据（含RMSD）已保存到: {group_avg_weight_csv_path}")
    
    def _perform_specific_comparison(self, analysis: TumorGrowthAnalysis, metric: str) -> None:
        """
        执行特定组之间的比较。
        """
        logger.info(f"\nPerforming a specific comparison for total {metric} AUC between Control and Drug_A_High:")
        comparison_result = analysis.compare_groups(
            metric='total_auc',
            group1='Control',
            group2='Drug_A_High'
        )
        logger.info(f"Mann-Whitney U test p-value: {comparison_result['p_value']:.4f}")


if __name__ == '__main__':
    # 运行肿瘤AUC分析命令
    command = TumorAUCCommand()
    command.run()
